# -*- coding: utf-8 -*-
"""
通用文件上传工具类
支持多种上传方式：本地复制、SFTP上传、API上传
具有良好的扩展性，可以轻松添加新的上传方式
"""

import os
import shutil
import logging
from abc import ABC, abstractmethod
from typing import Dict, List, Optional, Any, Callable
import requests
import paramiko
from pathlib import Path

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class UploadStrategy(ABC):
    """上传策略抽象基类"""
    
    @abstractmethod
    def upload_file(self, source_path: str, target_path: str, **kwargs) -> bool:
        """上传单个文件"""
        pass
    
    @abstractmethod
    def upload_files_by_prefix(self, source_dir: str, prefix_mapping: Dict[str, str], 
                              base_target_dir: str = "", **kwargs) -> Dict[str, List[str]]:
        """根据前缀映射批量上传文件"""
        pass


class LocalCopyStrategy(UploadStrategy):
    """本地文件复制策略"""
    
    def upload_file(self, source_path: str, target_path: str, **kwargs) -> bool:
        """复制单个文件到目标路径"""
        try:
            # 确保目标目录存在
            target_dir = os.path.dirname(target_path)
            os.makedirs(target_dir, exist_ok=True)
            
            # 复制文件
            shutil.copy2(source_path, target_path)
            logger.info(f"成功复制文件: {source_path} -> {target_path}")
            return True
        except Exception as e:
            logger.error(f"复制文件失败 {source_path} -> {target_path}: {e}")
            return False
    
    def upload_files_by_prefix(self, source_dir: str, prefix_mapping: Dict[str, str], 
                              base_target_dir: str = "", **kwargs) -> Dict[str, List[str]]:
        """根据前缀映射批量复制文件"""
        if not os.path.exists(source_dir):
            logger.error(f"源目录不存在: {source_dir}")
            return {"success": [], "failed": []}
        
        success_files = []
        failed_files = []
        
        for filename in os.listdir(source_dir):
            source_file_path = os.path.join(source_dir, filename)
            
            if not os.path.isfile(source_file_path):
                continue
                
            found_match = False
            for prefix, target_subdir in prefix_mapping.items():
                if filename.lower().startswith(prefix.lower()):
                    target_dir = os.path.join(base_target_dir, target_subdir)
                    target_file_path = os.path.join(target_dir, filename)
                    
                    if self.upload_file(source_file_path, target_file_path):
                        success_files.append(filename)
                    else:
                        failed_files.append(filename)
                    found_match = True
                    break
            
            if not found_match:
                logger.warning(f"文件 '{filename}' 不匹配任何前缀，跳过复制")
        
        return {"success": success_files, "failed": failed_files}


class SFTPUploadStrategy(UploadStrategy):
    """SFTP上传策略"""
    
    def __init__(self, host: str, port: int = 22, username: str = "", password: str = "", 
                 key_filename: str = None):
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.key_filename = key_filename
        self._sftp = None
        self._transport = None
    
    def _get_sftp_connection(self):
        """获取SFTP连接"""
        if self._sftp is None:
            self._transport = paramiko.Transport((self.host, self.port))
            if self.key_filename:
                private_key = paramiko.RSAKey.from_private_key_file(self.key_filename)
                self._transport.connect(username=self.username, pkey=private_key)
            else:
                self._transport.connect(username=self.username, password=self.password)
            self._sftp = paramiko.SFTPClient.from_transport(self._transport)
        return self._sftp
    
    def _mkdirs(self, sftp, path: str):
        """递归创建远程目录"""
        dirs = []
        while path not in ('/', ''):
            dirs.append(path)
            path = os.path.dirname(path)
        for d in reversed(dirs):
            try:
                sftp.stat(d)
            except FileNotFoundError:
                sftp.mkdir(d)
    
    def upload_file(self, source_path: str, target_path: str, **kwargs) -> bool:
        """上传单个文件到SFTP服务器"""
        try:
            sftp = self._get_sftp_connection()
            remote_dir = os.path.dirname(target_path)
            
            # 递归创建远程目录
            self._mkdirs(sftp, remote_dir)
            
            # 上传文件
            sftp.put(source_path, target_path)
            logger.info(f"SFTP上传成功: {source_path} -> {target_path}")
            return True
        except Exception as e:
            logger.error(f"SFTP上传失败 {source_path} -> {target_path}: {e}")
            return False
    
    def upload_files_by_prefix(self, source_dir: str, prefix_mapping: Dict[str, str], 
                              base_target_dir: str = "", **kwargs) -> Dict[str, List[str]]:
        """根据前缀映射批量上传文件到SFTP"""
        if not os.path.exists(source_dir):
            logger.error(f"源目录不存在: {source_dir}")
            return {"success": [], "failed": []}
        
        success_files = []
        failed_files = []
        
        for filename in os.listdir(source_dir):
            source_file_path = os.path.join(source_dir, filename)
            
            if not os.path.isfile(source_file_path):
                continue
                
            found_match = False
            for prefix, target_subdir in prefix_mapping.items():
                if filename.lower().startswith(prefix.lower()):
                    remote_dest_dir = os.path.join(base_target_dir, target_subdir).replace('\\', '/')
                    remote_file_path = os.path.join(remote_dest_dir, filename).replace('\\', '/')
                    
                    if self.upload_file(source_file_path, remote_file_path):
                        success_files.append(filename)
                    else:
                        failed_files.append(filename)
                    found_match = True
                    break
            
            if not found_match:
                logger.warning(f"文件 '{filename}' 不匹配任何前缀，跳过上传")
        
        return {"success": success_files, "failed": failed_files}
    
    def close(self):
        """关闭SFTP连接"""
        if self._sftp:
            self._sftp.close()
        if self._transport:
            self._transport.close()


class APIUploadStrategy(UploadStrategy):
    """API上传策略"""
    
    def __init__(self, base_url: str, api_key: str = "", headers: Dict[str, str] = None):
        self.base_url = base_url
        self.api_key = api_key
        self.headers = headers or {}
        if api_key:
            self.headers["x-api-key"] = api_key
    
    def upload_file(self, source_path: str, target_path: str, **kwargs) -> bool:
        """通过API上传单个文件"""
        try:
            # 从kwargs中获取queue_name，如果没有则从target_path推断
            queue_name = kwargs.get('queue_name', 'default_queue')
            organization_id = kwargs.get('organization_id', 'default_org')
            
            upload_url = f"{self.base_url}/{organization_id}/{queue_name}/in-resource/upload"
            
            with open(source_path, 'rb') as file:
                files = {'file': file}
                response = requests.post(
                    url=upload_url,
                    headers=self.headers,
                    files=files
                )
            
            if response.status_code == 200:
                logger.info(f"API上传成功: {source_path} -> {queue_name}, 响应: {response.json()}")
                return True
            else:
                logger.error(f"API上传失败: {source_path} -> {queue_name}, 状态码: {response.status_code}, 响应: {response.json()}")
                return False
                
        except Exception as e:
            logger.error(f"API上传异常 {source_path}: {e}")
            return False
    
    def upload_files_by_prefix(self, source_dir: str, prefix_mapping: Dict[str, str], 
                              base_target_dir: str = "", **kwargs) -> Dict[str, List[str]]:
        """根据前缀映射批量上传文件到API"""
        if not os.path.exists(source_dir):
            logger.error(f"源目录不存在: {source_dir}")
            return {"success": [], "failed": []}
        
        success_files = []
        failed_files = []
        
        for root, dirs, files in os.walk(source_dir):
            for filename in files:
                file_path = os.path.join(root, filename)
                lower_name = filename.lower()
                
                found_match = False
                for prefix, api_name in prefix_mapping.items():
                    if lower_name.startswith(prefix.lower()):
                        logger.info(f"[API上传] 调用队列: {api_name}, 上传文件: {file_path}")
                        
                        if self.upload_file(file_path, "", queue_name=api_name, **kwargs):
                            success_files.append(filename)
                        else:
                            failed_files.append(filename)
                        found_match = True
                        break
                
                if not found_match:
                    logger.warning(f"文件 '{filename}' 不匹配任何前缀，跳过API上传")
        
        return {"success": success_files, "failed": failed_files}


class FileUploadManager:
    """文件上传管理器"""
    
    def __init__(self, strategy: UploadStrategy):
        self.strategy = strategy
    
    def upload_file(self, source_path: str, target_path: str, **kwargs) -> bool:
        """上传单个文件"""
        return self.strategy.upload_file(source_path, target_path, **kwargs)
    
    def upload_files_by_prefix(self, source_dir: str, prefix_mapping: Dict[str, str], 
                              base_target_dir: str = "", **kwargs) -> Dict[str, List[str]]:
        """根据前缀映射批量上传文件"""
        return self.strategy.upload_files_by_prefix(source_dir, prefix_mapping, base_target_dir, **kwargs)
    
    def close(self):
        """关闭连接（如果策略支持）"""
        if hasattr(self.strategy, 'close'):
            self.strategy.close()


class UploadConfig:
    """上传配置类"""
    
    def __init__(self):
        self.prefix_to_api_mapping = {
            "aggregate_": "ppa_daily_load_aggregate",
            "exchange_rate_": "ppa_daily_load_exchange_rate",
            "portfolio_data_": "ppa_daily_load_portfolio",
            "portfolio_grouping_": "ppa_daily_load_portfolio_grouping",
            "position_": "ppa_daily_load_position",
            "transaction_": "ppa_daily_load_transaction",
            "benchmark_performance_": "ppa_monthly_load_benchmark_performance",
            "external_cash_flow_": "ppa_monthly_load_external_cash_flow",
            "portfolio_management_": "ppa_monthly_load_portfolio_management",
            "historical_pl_": "ppa_load_historical_pl",
            "historical_pl_attribution_": "ppa_load_historical_pl_attribution",
            "historical_returns_": "ppa_load_historical_returns"
        }
        
        self.prefix_to_folder_mapping = {
            'portfolio_data': 'portfolio',
            'aggregate': 'aggregate',
            'exchange_rate': 'exchange_rate',
            'portfolio_grouping': 'portfolio_grouping',
            'position': 'position',
            'transaction': 'transaction',
            'benchmark_performance': 'benchmark_performance',
            'external_cash_flow': 'external_cash_flow',
            'portfolio_management': 'portfolio_management',
            'portfolio_benchmark': 'portfolio_benchmark'
        }
    
    @classmethod
    def from_dict(cls, api_mapping: Dict[str, str] = None, folder_mapping: Dict[str, str] = None):
        """从字典创建配置"""
        config = cls()
        if api_mapping:
            config.prefix_to_api_mapping.update(api_mapping)
        if folder_mapping:
            config.prefix_to_folder_mapping.update(folder_mapping)
        return config


# 便捷函数
def create_local_uploader() -> FileUploadManager:
    """创建本地复制上传器"""
    return FileUploadManager(LocalCopyStrategy())


def create_sftp_uploader(host: str, port: int = 22, username: str = "", password: str = "", 
                        key_filename: str = None) -> FileUploadManager:
    """创建SFTP上传器"""
    strategy = SFTPUploadStrategy(host, port, username, password, key_filename)
    return FileUploadManager(strategy)


def create_api_uploader(base_url: str, api_key: str = "", headers: Dict[str, str] = None) -> FileUploadManager:
    """创建API上传器"""
    strategy = APIUploadStrategy(base_url, api_key, headers)
    return FileUploadManager(strategy)


# 使用示例
if __name__ == "__main__":
    # 示例1: 本地复制
    local_uploader = create_local_uploader()
    source_dir = r"D:\test_data"
    target_base_dir = r"E:\target_data"
    prefix_mapping = {
        'portfolio_data': 'portfolio',
        'aggregate': 'aggregate'
    }
    
    result = local_uploader.upload_files_by_prefix(source_dir, prefix_mapping, target_base_dir)
    print(f"本地复制结果: {result}")
    
    # 示例2: SFTP上传
    sftp_uploader = create_sftp_uploader(
        host='192.168.5.112',
        port=22,
        username='dockeruser01',
        password='zxcvbnm@158'
    )
    
    remote_base_dir = '/opt/application/ppa_qa/COMPASS_Workflow/CDIA-DEPLOYMENT/analytics-components/ppa_agent/simulation/source-files/hostname/in/data'
    result = sftp_uploader.upload_files_by_prefix(source_dir, prefix_mapping, remote_base_dir)
    print(f"SFTP上传结果: {result}")
    sftp_uploader.close()
    
    # 示例3: API上传
    api_uploader = create_api_uploader(
        base_url='http://192.168.5.112:9292/compass-workflow-api',
        api_key='eyJhbGciOiJIUzUxMiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJBeGlzb2Z0IENPTVBBU1MgQWxjaGVteUouSU8iLCJ1c2VySWQiOiJBTEFCL0JPQ0hLX1FBIiwianRpIjoiYzgxNjFiYTQtOGQ5Zi00MmJiLThlYjAtOGQ2ODc1Yzc5Y2MxIiwiZXhwIjozMjUwMzYyMjQwMH0.029SVgYcHJo9ibQVFajHGbof3OU1f2fV_N1Hjam5cJ4qaelv0zingnqulUYi9CHr2f9tXNBeDg6X-t9wWfCv2Q'
    )
    
    config = UploadConfig()
    result = api_uploader.upload_files_by_prefix(
        source_dir, 
        config.prefix_to_api_mapping, 
        organization_id='bochk'
    )
    print(f"API上传结果: {result}") 